"""
RAG摘要生成模块
使用检索增强生成技术生成论文摘要
"""
import logging
from typing import List, Dict, Optional
import requests
import json

from config import (
    USE_LOCAL_LLM, LOCAL_LLM_PATH,
    USE_OPENAI, OPENAI_API_KEY,
    SUMMARY_LENGTH, TEMPERATURE, MAX_NEW_TOKENS
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class RAGSummarizer:
    """RAG摘要生成器"""
    
    def __init__(self):
        self.use_local = USE_LOCAL_LLM
        self.use_openai = USE_OPENAI
        self.openai_key = OPENAI_API_KEY
        
        # 如果不是本地LLM也不使用OpenAI，则使用一个简单的模板
        if not self.use_local and not self.use_openai:
            logger.warning("未配置本地LLM或OpenAI，将使用模板生成摘要")
    
    def generate_summary(
        self, 
        papers: List[Dict], 
        query: str = ""
    ) -> str:
        """
        为检索到的论文生成摘要
        
        Args:
            papers: 论文列表
            query: 原始查询
            
        Returns:
            生成的摘要文本
        """
        if not papers:
            return "未找到相关论文。"
        
        # 构建上下文
        context = self._build_context(papers)
        
        # 生成提示
        prompt = self._build_prompt(context, query)
        
        # 调用LLM生成摘要
        if self.use_openai and self.openai_key:
            summary = self._generate_with_openai(prompt)
        elif self.use_local:
            summary = self._generate_with_local_llm(prompt)
        else:
            # 使用模板方法
            summary = self._generate_with_template(papers, query)
        
        return summary
    
    def _build_context(self, papers: List[Dict]) -> str:
        """构建上下文文本"""
        context_parts = []
        
        for i, paper in enumerate(papers, 1):
            title = paper.get("title", "No Title")
            abstract = paper.get("abstract", "")
            authors = ", ".join(paper.get("authors", [])[:3])  # 前3个作者
            journal = paper.get("journal", "")
            
            context_parts.append(
                f"[Paper {i}]\n"
                f"Title: {title}\n"
                f"Authors: {authors}\n"
                f"Journal: {journal}\n"
                f"Abstract: {abstract[:500]}...\n"
            )
        
        return "\n".join(context_parts)
    
    def _build_prompt(self, context: str, query: str) -> str:
        """构建提示词"""
        base_prompt = f"""You are a medical research assistant. Based on the following research papers, generate a concise summary of the key findings and insights.

{context}

Please summarize the key points from these papers in {SUMMARY_LENGTH}, focusing on:
1. Main research findings
2. Important implications for medical practice or research
3. Key recommendations or conclusions

Summary:"""
        
        if query:
            base_prompt = f"""You are a medical research assistant. A researcher is looking for information about: "{query}"

Based on the following research papers, generate a concise summary of the most relevant findings.

{context}

Please summarize the key points that are most relevant to the query "{query}" in {SUMMARY_LENGTH}, focusing on:
1. Main research findings related to the query
2. Important implications
3. Key recommendations

Summary:"""
        
        return base_prompt
    
    def _generate_with_openai(self, prompt: str) -> str:
        """使用OpenAI API生成摘要"""
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.openai_key}"
            }
            
            data = {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {"role": "system", "content": "You are a helpful medical research assistant."},
                    {"role": "user", "content": prompt}
                ],
                "temperature": TEMPERATURE,
                "max_tokens": MAX_NEW_TOKENS
            }
            
            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers=headers,
                json=data,
                timeout=60
            )
            
            response.raise_for_status()
            result = response.json()
            summary = result["choices"][0]["message"]["content"]
            
            logger.info("使用OpenAI API成功生成摘要")
            return summary
            
        except Exception as e:
            logger.error(f"OpenAI API调用失败: {e}")
            return self._generate_with_template(None, "")
    
    def _generate_with_local_llm(self, prompt: str) -> str:
        """使用本地LLM生成摘要"""
        try:
            # 这里需要根据实际的本地LLM接口进行适配
            # 例如使用llama.cpp, Ollama, 或其他本地推理框架
            
            logger.info("本地LLM功能需要根据具体模型进行配置")
            return self._generate_with_template(None, "")
            
        except Exception as e:
            logger.error(f"本地LLM调用失败: {e}")
            return self._generate_with_template(None, "")
    
    def _generate_with_template(
        self, 
        papers: Optional[List[Dict]], 
        query: str
    ) -> str:
        """使用模板方法生成摘要（fallback）"""
        if not papers:
            return "暂无摘要信息。"
        
        summary_parts = []
        
        # 收集关键信息
        key_topics = set()
        for paper in papers:
            # 从标题提取关键词
            title = paper.get("title", "")
            keywords = paper.get("keywords", [])
            key_topics.update(keywords)
        
        # 生成摘要
        summary_parts.append("Based on the retrieved research papers:")
        
        # 添加主要发现
        if papers:
            first_paper = papers[0]
            summary_parts.append(
                f"1. {first_paper.get('title', 'The studies')} discusses "
                f"{first_paper.get('abstract', '')[:150]}..."
            )
        
        # 添加相关主题
        if key_topics:
            summary_parts.append(
                f"2. Key topics include: {', '.join(list(key_topics)[:5])}"
            )
        
        # 添加建议
        summary_parts.append(
            "3. These findings contribute to advancing medical research "
            "and may have implications for clinical practice."
        )
        
        return " ".join(summary_parts)


if __name__ == "__main__":
    # 测试摘要生成
    summarizer = RAGSummarizer()
    
    test_papers = [
        {
            "title": "COVID-19 Vaccination and Long-term Immunity",
            "abstract": "This study examines the long-term immune response to COVID-19 vaccines...",
            "authors": ["John Doe", "Jane Smith"],
            "journal": "Nature Medicine"
        }
    ]
    
    summary = summarizer.generate_summary(test_papers, query="COVID-19 immunity")
    print(summary)
